import os

from WSModule import WSModule
from WSThread import WSThread
from WSJob import WSJob
from WSResult import WSResult
from WSOption import WSOption
from WSOptionGeneratorDict import WSOptionGeneratorDict

class Dafs(WSModule):
	log_path = os.getcwd() + '/results/dafs-result'
	options = {"threads": WSOption("threads", "Threads count, default 10", 10, False, ['-t', '--threads']),
			   "url": WSOption("url", "Target URL", "", True, ['-u', '--url']),
			   "msymbol": WSOption("msymbol", "Symbol of mask position in target URL (default №)",
				   "№", False, ['--msymbol']),
			   "method": WSOption("method", "Request method (default - HEAD)",
				   "HEAD", False, ['--method']),
			   "brute_type": WSOptionGeneratorDict("brute_type", "Type of brute",
				   "", True, ['--brute_type'], 'brute_options'),
			   }
	def __init__(self, kernel):
		super().__init__(kernel)
		for option_name in self.options:
			self.options[option_name].module = self
	def run(self):
		job = self.options["brute_type"].dict_generator.getJob()
		result = DafsResult()
		for i in range(int(self.options['threads'].value)):
			worker = Worker(job, result, self.options['url'].value, self.options['method'].value,
				self.options['msymbol'].value)
			worker.daemon = True
			worker.start()
		job.join()
		print(str(result))
	def help(self):
		return "Dafs help"
	def finished(self):
		return self.kernel.finished()

import requests, queue

class Worker(WSThread):
	q = None
	method = None
	url = None
	mask_symbol = None

	def __init__(self, job, result, url, method, mask_symbol):
		super().__init__(job, result)
		self.q = job
		self.method = method
		self.url = url
		self.mask_symbol = mask_symbol

	def run(self):
		sess = requests.Session()
		while True:
			try:
				word = self.q.get()
				url = self.url.replace(self.mask_symbol, word)
				resp = sess.head(url)\
				if self.method.lower() == 'head'\
				else sess.head(url)
				if resp.status_code != 404:
					self.result.put(url)
			except queue.Empty:
				self.running = False
				break
			finally:
				self.q.task_done()

class DafsResult(WSResult):
	urls = []
	def __str__(self):
		return '\n'.join(self.urls)
	def put(self, url):
		self.urls.append(url)